import os

import mindspore as ms
import numpy as np
import torch

from segan.models.discriminator import Discriminator as Discriminator_ms
from segan.models.generator import Generator as Generator_ms
from segan_pt.models.discriminator import Discriminator as Discriminator_pt
from segan_pt.models.generator import Generator as Generator_pt
from print_param import *

def pytorch_params(pt_dict):
    # par_dict = torch.load(pth_file, map_location='cpu')
    print("="*40)
    for name, param in pt_dict.items():
        print(name, param.numpy().shape)
        # pt_params[name] = param.numpy()
    # return pt_params

def check_d(pth_path, ckpt_path):
    inp = np.random.uniform(-1, 1, (16, 2, 16384)).astype(np.float32)
    # 注意做单元测试时，需要给Cell打训练或推理的标签
    ms_disc = Discriminator_ms(2, [64, 128, 256, 512, 1024], 31, [4] * 5, pool_type='none', pool_slen=16).set_train(False)
    pt_disc = Discriminator_pt(2, [64, 128, 256, 512, 1024], 31, [4] * 5, pool_type='none', pool_slen=16).eval()
    pt_disc.load_state_dict(torch.load(pth_path, map_location='cpu'))
    ms.load_checkpoint(ckpt_path, ms_disc)
    print("========= pt_disc conv1.weight ==========")
    print(pt_disc.enc_blocks[0].conv.weight.detach().numpy().reshape((-1,))[:10])
    print("========= ms_disc conv1.weight ==========")
    print(ms_disc.enc_blocks[0].conv.weight.data.asnumpy().reshape((-1,))[:10])
    pt_res, _ = pt_disc(torch.from_numpy(inp))
    ms_res = ms_disc(ms.Tensor(inp))
    print("========= pt_disc res ==========")
    print(pt_res)
    print("========= ms_disc res ==========")
    print(ms_res)
    print("diff", np.max(np.abs(pt_res.detach().numpy() - ms_res.asnumpy())))

def check_g(pth_path, ckpt_path):
    inp = np.random.uniform(-1, 1, (16, 1, 16384)).astype(np.float32)
    # 注意做单元测试时，需要给Cell打训练或推理的标签
    ms_genr = Generator_ms(1, [64, 128, 256, 512, 1024], kwidth=31, poolings=[4, 4, 4, 4, 4], no_z=True, bias=True, skip_merge='concat').set_train(False)
    pt_genr = Generator_pt(1, [64, 128, 256, 512, 1024], kwidth=31, poolings=[4, 4, 4, 4, 4], no_z=True, bias=True, skip_merge='concat').eval()
    # pytorch_params(pt_genr.state_dict())
    state_dict = torch.load(pth_path, map_location='cpu')
    # pytorch_params(state_dict)
    pt_genr.load_state_dict(state_dict)
    ms.load_checkpoint(ckpt_path, ms_genr)
    print("========= pt_genr conv1.weight ==========")
    print(pt_genr.enc_blocks[0].act.weight.detach().numpy().reshape((-1,))[:10])
    print("========= ms_genr conv1.weight ==========")
    print(ms_genr.enc_blocks[0].act.w.data.asnumpy().reshape((-1,))[:10])
    pt_res, pt_h_all = pt_genr(torch.from_numpy(inp), ret_hid = True)
    ms_res, ms_h_all = ms_genr(ms.Tensor(inp), ret_hid = True)
    for i, (pt_h, ms_h) in enumerate(zip(pt_h_all.values(), ms_h_all.values())):
        print("\n\n{:=^100s}".format(f" pt_genr res for layer {i} "))
        pt_h = pt_h.detach().numpy().reshape((-1,))
        print(pt_h[:20])
        print("\n{:=^100s}".format(f" ms_genr res for layer {i} "))
        ms_h = ms_h.asnumpy().reshape((-1,))
        print(ms_h[:20])
        print("diff", np.max(np.abs(pt_h - ms_h)))

ms.set_context(mode=ms.PYNATIVE_MODE)
os.environ['CUDA_VISIBLE_DEVICES'] = '7'

# pth_path = "/root/linx/segan_pt/pt_disciminator.pth"
# ckpt_path = "converted_d.ckpt"
# check_d(pth_path, ckpt_path)

pth_path = "/root/linx/segan_pt/pt_generator.pth"
ckpt_path = "converted_g.ckpt"
check_g(pth_path, ckpt_path)